{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "import opt_configs\n",
    "\n",
    "from opt_results import results\n",
    "\n",
    "from constants import OUTPUT_DIR\n",
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib.ticker import FormatStrFormatter\n",
    "\n",
    "import json\n",
    "\n",
    "sns.set()\n",
    "sns.set_style('darkgrid', {'axes.edgecolor':'black'})\n",
    "matplotlib.rcParams['axes.linewidth']= 0.5\n",
    "\n",
    "fig_name = 'opt_results'\n",
    "fig_dir = join(FIGURE_DIR, fig_name)\n",
    "used_configs = [configs.Warp()]\n",
    "names = ['Ours']\n",
    "\n",
    "show_iters = [0, 64, 128, 256, 'final']\n",
    "ref_views = [0, 4, 8, 10]\n",
    "\n",
    "fontsize = 12\n",
    "base_size = 2\n",
    "n_rows = len(results)\n",
    "n_cols = len(show_iters) + 3\n",
    "total_width = base_size * n_cols\n",
    "aspect = 0.98 * n_rows / n_cols\n",
    "\n",
    "total_width = TEXT_WIDTH\n",
    "fig = plt.figure(1, figsize=(total_width, aspect * total_width), constrained_layout=False)\n",
    "iter_rename_dict = {show_iters[-1]: 'Final', 0: 'Initialization'}\n",
    "gs = fig.add_gridspec(n_rows, n_cols, wspace=0.025, hspace=0.025)\n",
    "\n",
    "inset_size = 0.2\n",
    "y_offset = -0.25\n",
    "ranges = [5e-7, 2e-6, 5e-5]\n",
    "for row, result in enumerate(results):\n",
    "    scene, opt_name, pretty_name = result['scene'], result['opt_config'], result['pretty_name']\n",
    "    opt_config = opt_configs.get_opt_config(opt_name)\n",
    "    n_views = len(opt_config.sensors)\n",
    "    output_dir = join(OUTPUT_DIR, scene, opt_config.name, used_configs[0].name)\n",
    "\n",
    "    # ax = plt.Subplot(fig, gs[row, 0])\n",
    "    ax = fig.add_subplot(gs[row, 0])\n",
    "    ax.set_ylabel('\\\\textsc{' + pretty_name + '}' + '\\n' + f'{n_views} views', labelpad=5, fontsize=fontsize)\n",
    "    if row == n_rows - 1:\n",
    "        txt = ax.set_title('Reference images \\n (subset)', fontsize=fontsize, y=y_offset, va='top')\n",
    "    disable_ticks(ax)\n",
    "    disable_border(ax)\n",
    "\n",
    "    gs_insets = gridspec.GridSpecFromSubplotSpec(2, 2, subplot_spec=gs[row, 0], wspace=0.00, hspace=0.05)\n",
    "\n",
    "    for idx, ref_view in enumerate(ref_views):\n",
    "        r = idx // 2\n",
    "        c = idx % 2\n",
    "        ax = fig.add_subplot(gs_insets[r, c])\n",
    "        img = read_img(join(output_dir, f'ref-{ref_view:02d}.exr'))\n",
    "        ax.imshow(img, interpolation='none')\n",
    "        disable_ticks(ax)\n",
    "    col = 1\n",
    "    r = None\n",
    "    for cfg in used_configs:\n",
    "        for show_iter in show_iters:\n",
    "            ax = fig.add_subplot(gs[row, col])\n",
    "            suffix = f'{show_iter:04d}' if isinstance(show_iter, int) else show_iter\n",
    "            fn = join(FIGURE_DIR, fig_name, scene, f'{cfg.name}_{suffix}.exr')\n",
    "            img = read_img(fn)\n",
    "            img = img[90:-90, 90:-90, :]\n",
    "            ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "            disable_ticks(ax)\n",
    "            if row == n_rows - 1 and col > 0:\n",
    "                label = iter_rename_dict[show_iter] if show_iter in iter_rename_dict else f'{show_iter}'\n",
    "                txt = ax.set_xlabel(label, fontsize=fontsize, y=-0.5)\n",
    "                ax.xaxis.set_label_coords(0.5, -0.05)\n",
    "            if row == n_rows - 1 and col == len(show_iters) // 2 + 1:\n",
    "                txt = ax.set_title('Optimization states', fontsize=fontsize, y=y_offset, va='top')\n",
    "            col += 1\n",
    "\n",
    "        ax = fig.add_subplot(gs[row, col])\n",
    "        fn = join(FIGURE_DIR, fig_name, scene, f'reference.exr')\n",
    "        img = read_img(fn)\n",
    "        img = img[90:-90, 90:-90, :]\n",
    "\n",
    "        ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "        disable_ticks(ax)\n",
    "        if row == n_rows - 1:\n",
    "            txt = ax.set_title('Reference', fontsize=fontsize, y=y_offset, va='top')\n",
    "        col += 1\n",
    "\n",
    "        # Plot the convergence of the loss values\n",
    "        output_dir = join(OUTPUT_DIR, scene, opt_config.name, cfg.name)\n",
    "        with open(join(output_dir, f'metadata.json'), 'r') as f:\n",
    "            stats = json.load(f)\n",
    "        ax = fig.add_subplot(gs[row, col])\n",
    "        loss_scale = 100\n",
    "        loss_values = loss_scale * np.array(stats['loss_values'])\n",
    "        ax.plot(loss_values, color=sns.color_palette()[0], alpha=0.1)\n",
    "        ax.plot(smooth_loss(loss_values, 0.85), color=sns.color_palette()[0])\n",
    "\n",
    "        ax.yaxis.tick_right()\n",
    "        ax.yaxis.set_label_position(\"right\")\n",
    "        ax.set_ylabel(f'$L_1$ loss')\n",
    "        ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))\n",
    "        ax.set_ylim(0.01, np.quantile(loss_values, 0.99))\n",
    "        ax.set_xlim(0, 512)\n",
    "        ax.axes.get_xaxis().set_ticks([256, 512])\n",
    "        ax.tick_params(axis='x', pad=-4)\n",
    "        if row == n_rows - 1:\n",
    "            txt = ax.set_title('Iterations', fontsize=fontsize, y=y_offset, va='top')\n",
    "        else:\n",
    "            ax.axes.get_xaxis().set_ticklabels([])\n",
    "\n",
    "plt.margins(0, 0)\n",
    "# save_fig(fig_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
